##################################################
# (1) init ==> B <= 0
# (2) unsafe ==> B > 0
# (3) domain /\ x /\ x' ==> B(x') - (1-alpha) * B(x) < 0
##################################################
# max(B(x))  s.t. x in init
# if max(B)<=0: condition (1) is OK
import math
import time
import cvxpy as cp
import dccp
import gurobipy as gp
from gurobipy import GRB

import verify_dccp.getMathExpress2 as express
from prob import *
import numpy as np
import verify_piecewise_utils
from NN import cdinn2
import torch


def verif(model, piece_len):
    cex_init = []
    res_init = []
    cex_unsafe = []
    res_unsafe = []
    cex_domain = []
    res_domain = []

    x = cp.Variable(DIM)  # x0, x1, x2
    v = cp.Variable()

    f1, f2 = express.getExpr(x, model)

    # print("curvature of f1:", f1.curvature)
    # print("curvature of f2:", f2.curvature)

    # print("=====================X_INIT=========================")
    objective = cp.Maximize(f1 - v)
    constraints = [v == f2]

    constraints += [cp.square(x[0] - 1.5) + cp.square(x[1]) + cp.square(x[2]) <= 0.25 * 0.25]

    problem = cp.Problem(objective, constraints)
    # print("problem is DCCP:", dccp.is_dccp(problem))
    result = problem.solve(method='dccp')
    # print("x.value:", x.value)  # [ 1.34102711 -0.47405446]
    # print("v.value:", v.value)  # 4.158081062071981
    # print("status:", problem.status) # Converged
    # print("cost value0 =", result[0])  # -1.161325675958789
    if result[0] > 0.0:
        cex_init.append(x.value)
        res_init.append(result[0])
    else:
        print("init is safe")

    # print("=====================X_UNSAFE=========================")
    objective = cp.Minimize(v - f2)
    constraints = [v == f1]

    constraints += [cp.square(x[0] + 1) + cp.square(x[1] + 1) + cp.square(x[2] - 1) <= 0.64]

    problem = cp.Problem(objective, constraints)
    # print("problem is DCCP:", dccp.is_dccp(problem))
    # x.value = np.array([1, 2])
    result = problem.solve(method='dccp')
    # print("x.value:", x.value)
    # print("v.value:", v.value)
    # print("status:", problem.status)
    # print("cost value0 =", result[0])
    if result[0] < 0.0:
        cex_unsafe.append(x.value)
        res_unsafe.append(result[0])
    else:
        print("unsafe is safe")

    # print("=====================X_DOMAIN=========================")

    # verify_c8_1.txt -3_3_1_20.txt
    # cex_domain: [array([-3., -1., -2.26343815]), array([-3.00000000e+00, 1.32681603e-12, -2.44612817e+00]),
    #              array([-3., 0.87613818, -2.04550256]), array([-2.85714286, 1., -2.])]
    # res_domain: [0.24000771943074395, 1.4298883853902846, 2.215873472987629, 2.110507592438094]

    # fw = open('./log/verify_c8_1.txt', 'w+', encoding='utf-8')

    data_dir = "./data/"
    pieces = int(6 / piece_len)
    datas = np.loadtxt(data_dir + f"-3_3_{piece_len}_20.txt", dtype=np.float64).reshape([-1, 4])
    new_cex_idx_list = []

    start_time_cond3 = time.time()

    var_b = cp.Variable()

    x0_min = domain_min[0]  # -3
    x0_max = domain_max[0]  # 3
    x1_min = domain_min[1]  # -3
    x1_max = domain_max[1]  # 3
    x2_min = domain_min[2]  # -3
    x2_max = domain_max[2]  # 3
    # print(x0_min)
    # print(x1_min)
    # print(x2_min)
    # print(x0_max)
    # print(x1_max)
    # print(x2_max)

    # cex_idx_list = [[4, 2, 2], [4, 2, 3], [4, 3, 2], [4, 3, 3], [5, 2, 2], [5, 2, 3], [5, 3, 2], [5, 3, 3]]
    cex_idx_list = []

    if len(cex_idx_list) == 0:
        for x0_idx in range(pieces):
            for x1_idx in range(pieces):
                for x2_idx in range(pieces):
                    cex_idx = [x0_idx, x1_idx, x2_idx]
                    cex_idx_list.append(cex_idx)

    # print(cex_idx_list)

    for cex_idx in cex_idx_list:
        x0_idx, x1_idx, x2_idx = cex_idx

        # w0, w1, b_l, b = datas[x0_idx * pieces + x2_idx - 1]
        w0, w1, b_l, b = datas[x0_idx * pieces + x2_idx]

        x0_l = x0_idx * piece_len + x0_min
        x0_r = (x0_idx + 1) * piece_len + x0_min
        x1_l = x1_idx * piece_len + x1_min
        x1_r = (x1_idx + 1) * piece_len + x1_min
        x2_l = x2_idx * piece_len + x2_min
        x2_r = (x2_idx + 1) * piece_len + x2_min
        # print(x0_l)
        # print(x0_r)
        # print(x1_l)
        # print(x1_r)
        # print(x2_l)
        # print(x2_r)
        # print("============")

        x0_next = time_step * (- x[0] + x[1] - x[2]) + x[0]
        x1_next = time_step * (- x[0] - x[1] - (w0 * x[0] + w1 * x[2] + b_l + var_b)) + x[1]
        x2_next = time_step * (- x[0] + 1.77 * x[0] - 4.8 * x[2]) + x[2]

        f1_next, f2_next = express.getExpr([x0_next, x1_next, x2_next], model)

        # print("curvature of f1_next:", f1_next.curvature)
        # print("curvature of f2_next:", f2_next.curvature)

        u = cp.Variable()
        objective = cp.Maximize((f1_next - v) - (1 - alpha) * (u - f2))
        constraints = [v == f2_next, u == f1]
        constraints += [x0_l <= x[0], x[0] <= x0_r]
        constraints += [x1_l <= x[1], x[1] <= x1_r]
        constraints += [x2_l <= x[2], x[2] <= x2_r]
        constraints += [x0_min <= x0_next, x0_next <= x0_max]
        constraints += [x1_min <= x1_next, x1_next <= x1_max]
        constraints += [x2_min <= x2_next, x2_next <= x2_max]

        constraints += [0 <= var_b, var_b <= b]

        problem = cp.Problem(objective, constraints)
        # print("problem is DCCP:", dccp.is_dccp(problem))
        # x.value = np.array([1, 2])
        result = problem.solve(method='dccp')
        # print("x.value:", x.value)
        # print("var_b.value:", var_b.value)
        # print("problem.status:", problem.status)
        if problem.status == "Converged":
            # print("cost value =", result[0])
            # fw.write("================[" + str(x_l) + "," + str(x_r) + "]===================")
            # fw.write('\n')
            # fw.write(str(result[0]))
            # fw.write('\n')
            # fw.write('\n')

            if result[0] > 0.0:
                cex_domain.append(x.value)
                res_domain.append(result[0])
                # print(f"[{x0_l}, {x0_r}], [{x1_l}, {x1_r}], [{x2_l}, {x2_r}] has cex point")
                # fw.write(f"[{x0_l}, {x0_r}], [{x1_l}, {x1_r}], [{x2_l}, {x2_r}] has cex point\n")
                # new_cex_idx_list.append(cex_idx)
        #     else:
        #         print(f"[{x0_l}, {x0_r}], [{x1_l}, {x1_r}], [{x2_l}, {x2_r}] is safe")
        #         # fw.write(f"[{x0_l}, {x0_r}], [{x1_l}, {x1_r}], [{x2_l}, {x2_r}] is safe\n")
        # else:
        #     print(f"[{x0_l}, {x0_r}], [{x1_l}, {x1_r}], [{x2_l}, {x2_r}] is not converged")
        #     # fw.write(f"[{x0_l}, {x0_r}], [{x1_l}, {x1_r}], [{x2_l}, {x2_r}] is not converged\n")

    end_time_cond3 = time.time()
    # print(new_cex_idx_list)
    # fw.write(f"new_cex_idx_list = {new_cex_idx_list}\n")
    # fw.write(f"verify_cond3_time = {end_time_cond3 - start_time_cond3}\n")
    # fw.close()
    # print("==============cex and res=============")
    print("cex_init: ", cex_init)
    print("res_init: ", res_init)
    print("cex_unsafe: ", cex_unsafe)
    print("res_unsafe: ", res_unsafe)
    print("cex_domain: ", cex_domain)
    print("res_domain: ", res_domain)
    return cex_init, cex_unsafe, cex_domain


if __name__ == "__main__":
    model = cdinn2.gen_nn()
    model.load_state_dict(torch.load('../model/cdinn_c8_train.pt'), strict=True)
    # time_start_verify = time.time()
    verif(model)
    # time_end_verify = time.time()
    # print("verifing totally costs:", time_end_verify - time_start_verify)



